import os
import sys
from typing import List

import numpy as np
import random
from torch.utils.data import Dataset
from PIL import Image
import cv2

np.random.seed(0)

from torchvision import transforms  # noqa

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from configuration import Configuration

class ImageDataset(Dataset):

    def __init__(self, args:Configuration, is_valid_data:bool):
        self._is_valid_data = is_valid_data
        self._args = args
        # self._color_jitter = transforms.ColorJitter(64.0/255, 0.75, 0.25, 0.04)
        self._color_jitter = transforms.ColorJitter(
                brightness=args.color_jitter_brightness, 
                contrast=args.color_jitter_contrast, 
                saturation=args.color_jitter_saturation, 
                hue=args.color_jitter_hue)
        self._items = []
        for class_idx, class_name in enumerate(args.classes):
            if is_valid_data:
                list_file = os.path.join(args.base_dir, class_name + ".valid.list")
            else:
                list_file = os.path.join(args.base_dir, class_name + ".train.list")

            if not os.path.isfile(list_file):
                raise Exception("not find list file: " + list_file)

            with open(list_file) as f:
                self._items += [(line.strip(), class_idx) for line in f]

        # random.shuffle(self._items)

    def __len__(self):
        return len(self._items)

    def __getitem__(self, idx):
        path, label = self._items[idx]
        label = np.array(label, dtype=float)

        img = Image.open(os.path.join(self._args.patch_dir, path))

        if not self._is_valid_data:
            # color jitter
            img = self._color_jitter(img)

        img = cv2.cvtColor(np.asarray(img),cv2.COLOR_RGB2BGR)

        if not self._is_valid_data:
            # use left_right flip
            if np.random.rand() > 0.5:
                img = cv2.flip(img, 1)

            # use rotate
            num_rotate = np.random.randint(0, 4)
            if num_rotate == 0:
                pass
            elif num_rotate == 1:
                img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
            elif num_rotate == 2:
                img = cv2.rotate(img, cv2.ROTATE_180)
            else:
                img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)

        if self._args.hsv:
            img = cv2.cvtColor(img,cv2.COLOR_BGR2HSV)

        img = np.array(img, dtype=np.float32)

        # PIL image: H W C
        # torch image: C H W
        img = img.transpose((2, 0, 1))

        if self._args.normalize:
            img = (img - 128.0) / 128.0

        return img, label

